import time
import os.path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from astar import AStar
from graph import Vertex, Edge
from qtree import Rectangle, Quadtree, Point

def visualize_quadtree_map(grid, quadtree, path, path2, dx, dy):
    """可视化网格地图和四叉树分割"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))

    # 原始网格地图
    ax1.imshow(grid, cmap=ListedColormap(['black', 'white']), origin='lower')
    ax1.set_title("Original Grid Map")
    ax1.set_xlim(0, dx)
    ax1.set_ylim(0, dy)

    # 四叉树分割结果
    ax2.set_title("Quadtree Partitioning")
    quadtree.draw(ax2)
    ax2.set_xlim(0, dx)
    ax2.set_ylim(0, dy)
    ax2.set_aspect('equal')

    if path:
        # 绘制路径
        path_x = [node.boundary.x + node.boundary.width / 2 for node in path]
        path_y = [node.boundary.y + node.boundary.height / 2 for node in path]
        ax2.plot(path_x, path_y, 'r-', linewidth=2, label='A* Path')

    if path2:
        # 绘制路径
        path_x = [node.boundary.x + node.boundary.width / 2 for node in path2]
        path_y = [node.boundary.y + node.boundary.height / 2 for node in path2]
        ax2.plot(path_x, path_y, 'b-', linewidth=1, label='spark Path')

    plt.tight_layout()
    plt.show()

# 起点和终点（转换为四叉树叶节点）
def find_leaf_node(quadtree, x, y):
    if not quadtree.divided and quadtree.boundary.contains(Point(x, y)):
        return quadtree
    for child in quadtree.children:
        node = find_leaf_node(child, x, y)
        if node:
            return node
    return None

# 提取四叉树中所有可通行的叶子节点作为顶点
def extract_passable_leaves(quadtree_root):
    leaves = []
    if quadtree_root.is_passable():
        leaves.append(quadtree_root)
    elif quadtree_root.divided:
        for child in quadtree_root.children:
            leaves.extend(extract_passable_leaves(child))
    return leaves

def test_quadtree_to_graph(qt):
    # 测试2，将四叉树数据转为图数据
    if not qt.is_buffer_neighbors:
        qt.build_neighbor_cache(qt)  # 构建边要知道一个节点的所有邻居。
    tt1 = time.perf_counter()
    l = extract_passable_leaves(qt)
    vertices = {}  # {(x0, y0): Vertex0, (x1, y1): Vertex1, ...} 这样组织可以方便我们下面根据坐标查邻居顶点。
    for node_n in l:
        vertices[(node_n.boundary.x, node_n.boundary.y)] = Vertex(node_n)
    edges = []
    for _, vertex in vertices.items():
        neighbors = vertex.qtree_node.get_neighbors(qt)
        for neb in neighbors:
            vertex_neb = vertices.get((neb.boundary.x, neb.boundary.y))  # dict of {(x,y), vertex}
            if vertex_neb is not None:
                edges.append(Edge(vertex, vertex_neb))
    tt2 = time.perf_counter()
    print("建立图的顶点和边，用时:", round(tt2 - tt1, 3), "秒")
    print("图中，顶点数：", len(vertices), "，边数：", len(edges))
    #################################################################
    # 把顶点和边以json格式各自保存到一个文件，文件中一行一个对象好方便spark读取。
    # v = {"id":0, "center_x": "0.005", "center_y": "0.001"}
    # 顶点ID等于： DX*x + y
    lines = []
    for k, v in vertices.items():
        ss = '{"id":' + str(int(DX * k[0] + k[1])) + ', "center_x": ' + str(v.center[0]) + ', "center_y": ' + str(
            v.center[1]) + '}'
        lines.append(ss)
    context = '\n'.join(lines)
    with open("testv.json", 'wb') as f:
        f.write(context.encode())
        f.close()
    lines.clear()

    for ed in edges:
        lid1, lid2 = int(DX * ed.v1.id[0] + ed.v1.id[1]), int(DX * ed.v2.id[0] + ed.v2.id[1])
        ss = '{"v1":' + str(lid1) + ', "v2":' + str(lid2) + ', "w":' + str(ed.weight) + '}'
        lines.append(ss)
    context = '\n'.join(lines)
    with open("teste.json", 'wb') as f:
        f.write(context.encode())
        f.close()


def test_grid_data():
    # d_lon = round((140.0 - 100.0) / 0.01) # 4000
    # d_lat = round((23.634 - (-6.366)) / 0.01) # 3000
    DX, DY = 4000, 3000
    # 生成地图
    # (1) 使用spark生成的二进制文件作为数据源
    with open("D:\\temp\\jl\\grids.data", "rb") as f:
        grid = np.fromfile(f, dtype=np.byte, count=DX * DY)
    # 1:     1, 2, 3, ...., 3999,4000
    # 2:     1, 2, 3, ...., 3999,4000
    # 3:     1, 2, 3, ...., 3999,4000
    # ......
    # 3000:  1, 2, 3, ...., 3999,4000
    # 共3000行,4000列。对np来说，shape是(3000, 4000)
    grid = grid.reshape(DY, DX)
    # (2) 用python脚本grid_maker生成的hbase数据再次dump出来的缓存文件。
    # grid = np.load("grid_bufXXX.npy")
    tm1 = time.perf_counter()
    # 构建四叉树
    boundary = Rectangle(0, 0, DX, DY)
    qt = Quadtree(boundary, grid, max_depth=9)
    tm2 = time.perf_counter()
    print("建立四叉树用时:", round(tm2 - tm1, 3), "秒")

    # 取消下方的注释可以建立缓存
    t0 = time.perf_counter()
    qt.build_neighbor_cache(qt)
    t1 = time.perf_counter()
    print("建立缓存用时:", round(t1-t0, 3), "秒")

    start_node = find_leaf_node(qt, 110, 339)  # 起点 (10,10)
    goal_node = find_leaf_node(qt, 485, 375)  # 终点 (240,240)
    print("real src is ({}, {}) , dst is ({}, {}).".format(start_node.boundary.x,
                                                           start_node.boundary.y,
                                                           goal_node.boundary.x,
                                                           goal_node.boundary.y))
    start_node.test = goal_node.test = {"mask": True, "color": 'red'}
    # 测试1，展示起点和终点周边的邻居
    # for i in start_node.get_neighbors(qt):
    #     i.test["mask"] = True
    #     i.test["color"] = 'blue'
    # for i in goal_node.get_neighbors(qt):
    #     i.test["mask"] = True
    #     i.test["color"] = 'blue'
    ##################################################################
    # 测试寻路
    path = []
    #path = AStar.find_path(start_node, goal_node, qt)
    # # nodes.txt中的内容是在其他程序中用图计算出来的节点列表。
    path2 = []
    # with open("nodes.txt", 'r') as f:
    #     s = f.read()
    #     temp_p = s.split(",")
    #     for i in temp_p:
    #         j = int(i)
    #         x, y = j / DX, j % DX
    #         node = find_leaf_node(qt, x, y)
    #         if node:
    #             path2.append(node)
    #     f.close()

    # 可视化
    visualize_quadtree_map(grid, qt, path, path2, DX, DY)

def test_512_npz():
    npz_binary_file = os.path.abspath(os.path.normpath("./data/tile/0_0.npy"))
    grid = np.load(npz_binary_file)
    print(type(grid))
    DX, DY = 512, 512

    tm1 = time.perf_counter()
    # 构建四叉树
    boundary = Rectangle(0, 0, DX, DY)
    qt = Quadtree(boundary, grid, max_depth=9)
    tm2 = time.perf_counter()
    print("建立四叉树用时:", round(tm2 - tm1, 3), "秒")

    # 取消下方的注释可以建立缓存
    t0 = time.perf_counter()
    qt.build_neighbor_cache(qt)
    t1 = time.perf_counter()
    print("建立缓存用时:", round(t1-t0, 3), "秒")

    start_node = find_leaf_node(qt, 0, 0)  # 起点 (10,10)
    goal_node = find_leaf_node(qt, 7, 510)  # 终点 (240,240)
    print("real src is ({}, {}) , dst is ({}, {}).".format(start_node.boundary.x,
                                                           start_node.boundary.y,
                                                           goal_node.boundary.x,
                                                           goal_node.boundary.y))
    start_node.test = goal_node.test = {"mask": True, "color": 'red'}
    # 测试寻路
    # path = []
    path = AStar.find_path(start_node, goal_node, qt)
    # nodes.txt中的内容是在其他程序中用图计算出来的节点列表。
    path2 = []
    # 可视化
    visualize_quadtree_map(grid, qt, path, path2, DX, DY)

if __name__ == '__main__':
    test_512_npz()